import base64
import io
from ultralytics import YOLO
from flask import request, Flask, jsonify
from waitress import serve
from PIL import Image
from flask_socketio import SocketIO

app = Flask(__name__)
model = YOLO("yolov8n.pt")

app.config['SECRET_KEY'] = 'secret!'
socketio = SocketIO(app)


@app.route("/socket")
def socket():
    with open("socket.html",encoding='utf-8') as file:
        return file.read()
    
@app.route("/")
def root():
    with open("index.html",encoding='utf-8') as file:
        return file.read()

@socketio.on('message')
def handle_message(data):
    base64_data = data.split(',')[1]
    image_bytes = base64.b64decode(base64_data)
    image = Image.open(io.BytesIO(image_bytes))
    json = detect_objects_on_image(image)
    socketio.send(json)

@app.route("/detect", methods=["POST"])
def detect():
    buf = request.files["image_file"]
    boxes = detect_objects_on_image(Image.open(buf.stream))
    return jsonify(boxes)

# 对象检测
def detect_objects_on_image(stream):
    results = model.predict(stream)
    result = results[0]
    output = []
    for box in result.boxes:
        x1, y1, x2, y2 = [
            round(x) for x in box.xyxy[0].tolist()
        ]
        class_id = box.cls[0].item()
        prob = round(box.conf[0].item(), 2)
        output.append([
            x1, y1, x2, y2, result.names[class_id], prob
        ])
    return output


if __name__ == '__main__':
    # 开启WebSocket支持
    socketio.run(app, debug=True)
    serve(app, host='0.0.0.0', port=8080)